# MicroPython SSD1306 OLED driver, I2C and SPI interfaces

from micropython import const
import framebuf


# register definitions
SET_CONTRAST        = const(0x81)
SET_ENTIRE_ON       = const(0xa4)
SET_NORM_INV        = const(0xa6)
SET_DISP            = const(0xae)
SET_MEM_ADDR        = const(0x20)
SET_COL_ADDR        = const(0x21)
SET_PAGE_ADDR       = const(0x22)
SET_DISP_START_LINE = const(0x40)
SET_SEG_REMAP       = const(0xa0)
SET_MUX_RATIO       = const(0xa8)
SET_COM_OUT_DIR     = const(0xc0)
SET_DISP_OFFSET     = const(0xd3)
SET_COM_PIN_CFG     = const(0xda)
SET_DISP_CLK_DIV    = const(0xd5)
SET_PRECHARGE       = const(0xd9)
SET_VCOM_DESEL      = const(0xdb)
SET_CHARGE_PUMP     = const(0x8d)

# Subclassing FrameBuffer provides support for graphics primitives
# http://docs.micropython.org/en/latest/pyboard/library/framebuf.html.


class SSD1306(framebuf.FrameBuffer):
	def __init__(self, width, height, external_vcc):
		self.width = width
		self.height = height
		self.external_vcc = external_vcc
		self.pages = self.height // 8
		self.buffer = bytearray(self.pages * self.width)
		super().__init__(self.buffer, self.width, self.height, framebuf.MONO_VLSB)
		self.init_display()

	def init_display(self):
		for cmd in (
			SET_DISP | 0x00,  # off
			# address setting
			SET_MEM_ADDR, 0x00,  # horizontal
			# resolution and layout
			SET_DISP_START_LINE | 0x00,
			SET_SEG_REMAP | 0x01,  # column addr 127 mapped to SEG0
			SET_MUX_RATIO, self.height - 1,
			SET_COM_OUT_DIR | 0x08,  # scan from COM[N] to COM0
			SET_DISP_OFFSET, 0x00,
			SET_COM_PIN_CFG, 0x02 if self.height == 32 else 0x12,
			# timing and driving scheme
			SET_DISP_CLK_DIV, 0x80,
			SET_PRECHARGE, 0x22 if self.external_vcc else 0xf1,
			SET_VCOM_DESEL, 0x30,  # 0.83*Vcc
			# display
			SET_CONTRAST, 0xff,  # maximum
			SET_ENTIRE_ON,  # output follows RAM contents
			SET_NORM_INV,  # not inverted
			# charge pump
			SET_CHARGE_PUMP, 0x10 if self.external_vcc else 0x14,
			SET_DISP | 0x01):  # on
			self.write_cmd(cmd)
		self.fill(0)
		self.show()

	def poweroff(self):
		self.write_cmd(SET_DISP | 0x00)

	def poweron(self):
		self.write_cmd(SET_DISP | 0x01)

	def contrast(self, contrast):
		self.write_cmd(SET_CONTRAST)
		self.write_cmd(contrast)

	def invert(self, invert):
		self.write_cmd(SET_NORM_INV | (invert & 1))

	def show(self):
		x0 = 0
		x1 = self.width - 1
		if self.width == 64:
			# displays with width of 64 pixels are shifted by 32
			x0 += 32
			x1 += 32
		self.write_cmd(SET_COL_ADDR)
		self.write_cmd(x0)
		self.write_cmd(x1)
		self.write_cmd(SET_PAGE_ADDR)
		self.write_cmd(0)
		self.write_cmd(self.pages - 1)
		self.write_data(self.buffer)

	def show_symbol(self, array, offset_x=0, offset_y=0):
		if not isinstance(array, list):
			raise TypeError("Wrong type of data was given")

		self.fill(0)

		for y, row in enumerate(array):
			for x, c in enumerate(row):
				self.pixel(offset_x + x, offset_y + y, c)

	def show_chinese(self, words: str, width, height, offset_x, offset_y, ignore=True) -> None:
		"""
		显示汉字字模
		:param words: 要显示的文字（中英文）
		:param width: 字模高度
		:param height: 字模宽度
		:param offset_x: 屏幕输出位置 X 轴偏移量
		:param offset_y: 屏幕输出位置 Y 轴偏移量
		:param ignore: 是否忽略错误
		"""

		from drivers.font import hanzi
		import math

		word_offset = 0  # 多个字之间的偏移量

		# 枚举所有汉字
		for word in words:
			word_code = 0x0  # 文字的 Unicode 编码

			code = word.encode("utf-8")

			# 根据字节长度判断中英文
			if len(code) == 3:
				word_code |= code[0] << 16
				word_code |= code[1] << 8
				word_code |= code[2]
			else:
				word_code |= code[0]

			try:
				# 从字体库读取字模，如果不存在则给出错误提示
				matrix_array = hanzi.hanzi[word_code]
			except KeyError:
				if ignore:
					print("Missing matrix define: '{0}'\n".format(word))

					continue
				else:
					raise KeyError("Missing matrix define: ", word)

			matrix_index = 0  # 字模索引下标

			# 按行显示字模
			for count in range(0, height):
				row_count = 0  # 每行打印点阵数量统计

				# 根据每行宽度计算打印次数，也就是每行需要从字模列表里取几次数据
				for times in range(0, math.ceil(width / 8)):
					# 将取到的字模数据转换为二进制字符串（8 位 0 和 1 的组合）
					# 因为字模宽度不一定为 8 的整倍数，所以在取每行最后一个数据时，
					# 需要根据实际情况获取指定位数的二进制字符串
					binary = bin(matrix_array[matrix_index]).replace('0b', '')
					binary = self.rjust(binary, width % 8 if times + 1 == math.ceil(width / 8) else 8)
					matrix_index += 1

					# print(binary, end='')

					# 顺序打印每行点阵
					for bit in binary:
						# print("●" if bit == '1' else "○", end='')

						if bit == '1':
							self.pixel(word_offset + offset_x + row_count, offset_y + count, 1)

						row_count += 1

						# 如果已打印的点阵数量达到指定宽度，
						# 则跳出循环并开始打印下一行
						if row_count > width:
							break
			# 	print()
			# print()

			word_offset += width

	@staticmethod
	def rjust(bits: str, count: int):
		count = 8 if count == 0 else count

		for i in range(0, count):
			bits = '0' + bits

		bits = bits[count - count * 2:]

		return bits


class SSD1306_I2C(SSD1306):
	def __init__(self, width, height, i2c, addr=0x3c, external_vcc=False):
		self.i2c = i2c
		self.addr = addr
		self.temp = bytearray(2)
		self.write_list = [b'\x40', None] # Co=0, D/C#=1
		super().__init__(width, height, external_vcc)

	def write_cmd(self, cmd):
		self.temp[0] = 0x80 # Co=1, D/C#=0
		self.temp[1] = cmd
		self.i2c.writeto(self.addr, self.temp)

	def write_data(self, buf):
		self.write_list[1] = buf
		self.i2c.writevto(self.addr, self.write_list)


class SSD1306_SPI(SSD1306):
	def __init__(self, width, height, spi, dc, res, cs, external_vcc=False):
		self.rate = 10 * 1024 * 1024
		dc.init(dc.OUT, value=0)
		res.init(res.OUT, value=0)
		cs.init(cs.OUT, value=1)
		self.spi = spi
		self.dc = dc
		self.res = res
		self.cs = cs
		import time
		self.res(1)
		time.sleep_ms(1)
		self.res(0)
		time.sleep_ms(10)
		self.res(1)
		super().__init__(width, height, external_vcc)

	def write_cmd(self, cmd):
		self.spi.init(baudrate=self.rate, polarity=0, phase=0)
		self.cs(1)
		self.dc(0)
		self.cs(0)
		self.spi.write(bytearray([cmd]))
		self.cs(1)

	def write_data(self, buf):
		self.spi.init(baudrate=self.rate, polarity=0, phase=0)
		self.cs(1)
		self.dc(1)
		self.cs(0)
		self.spi.write(buf)
		self.cs(1)
